#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
AI营销方案生成系统 - WebSocket处理模块
提供实时通信功能，包括AI对话流式传输

主要功能：
- WebSocket连接管理
- AI对话流式传输
- 实时消息推送
- 连接状态管理

Author: AI Marketing System
Date: 2024
"""

import json
import logging
from typing import Dict, Any, Optional
from flask import request, session
from flask_socketio import SocketIO, emit, join_room, leave_room, disconnect
from functools import wraps

# 日志配置
logger = logging.getLogger(__name__)

# WebSocket处理器注册函数
def register_websocket_handlers(socketio: SocketIO):
    """
    注册WebSocket事件处理器到指定的SocketIO实例
    
    Args:
        socketio: SocketIO实例
    """

    @socketio.on('connect')
    def handle_connect(auth=None):
        """
        处理WebSocket连接
        
        Args:
            auth: 认证信息（可选）
        """
        client_id = request.sid
        logger.info(f"WebSocket客户端连接 - ID: {client_id}")
        
        # 发送连接成功消息
        emit('connect_success', {
            'message': '连接成功',
            'client_id': client_id
        })
    
    @socketio.on('disconnect')
    def handle_disconnect():
        """
        处理WebSocket断开连接
        """
        client_id = request.sid
        logger.info(f"WebSocket客户端断开连接 - ID: {client_id}")
    
    @socketio.on('start_generation')
    def handle_start_generation(data):
        """
        处理AI流式生成请求
        
        Args:
            data: 包含requirement_id, session_id, message的字典
        """
        try:
            client_id = request.sid
            
            # 参数验证
            if not data:
                emit('stream_error', {'error': '缺少请求数据'})
                return
            
            requirement_id = data.get('requirement_id')
            session_id = data.get('session_id')
            message = data.get('message', '')
            
            if not requirement_id or not session_id:
                emit('stream_error', {'error': '缺少必要参数: requirement_id 或 session_id'})
                return
            
            logger.info(f"开始AI流式生成 - 客户端: {client_id}, 会话: {session_id}")
            
            # 发送开始生成事件
            emit('stream_message', {
                'type': 'start',
                'message': '开始生成营销方案...'
            })
            
            # 获取需求信息用于构建AI提示词
            requirement_data = None
            if requirement_id:
                from services.requirement_service import RequirementService
                requirement_service = RequirementService()
                requirement_data = requirement_service.get_requirement(requirement_id)
            
            # 构建AI提示词
            prompt = _build_marketing_prompt(requirement_data, data.get('message', ''))
            logger.info(f"构建AI提示词完成 - 会话: {session_id}")
            
            # 初始化AI服务
            from services.ai_service import AIService
            ai_service = AIService()
            
            # 使用AI服务进行流式生成
            full_content = ""
            chunk_count = 0
            
            for chunk in ai_service.generate_content_stream(prompt):
                if chunk:
                    full_content += chunk
                    chunk_count += 1
                    
                    # 发送流式数据块
                    emit('stream_message', {
                        'type': 'chunk',
                        'content': chunk,
                        'progress': min(chunk_count * 2, 95)  # 动态进度，最大95%
                    })
            
            # 发送完成事件
            emit('stream_message', {
                'type': 'complete',
                'message': '营销方案生成完成',
                'full_content': full_content
            })
            
            logger.info(f"AI流式生成完成 - 客户端: {client_id}, 会话: {session_id}")
            
        except Exception as e:
            logger.error(f"AI生成过程中发生错误: {str(e)}")
            emit('stream_error', {
                'error': f'生成过程中发生错误: {str(e)}'
            })
    
    @socketio.on('stop_generation')
    def handle_stop_generation(data=None):
        """
        处理停止生成请求
        
        Args:
            data: 请求数据（可选）
        """
        client_id = request.sid
        logger.info(f"停止AI生成 - 客户端: {client_id}")
        
        emit('stream_message', {
            'type': 'stopped',
            'message': '生成已停止'
        })
    
    @socketio.on('ping')
    def handle_ping(data=None):
        """
        处理心跳检测
        
        Args:
            data: 心跳数据（可选）
        """
        import time
        client_id = request.sid
        logger.debug(f"收到心跳 - 客户端: {client_id}")
        
        emit('pong', {
            'message': 'pong',
            'timestamp': int(time.time() * 1000)
        })
    
    @staticmethod
    def _build_marketing_prompt(requirement_data, user_message=""):
        """
        构建营销方案生成的AI提示词
        
        Args:
            requirement_data: 需求数据
            user_message: 用户消息
            
        Returns:
            str: 构建好的提示词
        """
        base_prompt = """
你是一位专业的营销策划专家，请根据以下需求信息生成一份详细的营销方案。

请按照以下结构生成方案：
1. 项目概述
2. 目标分析
3. 客户偏好洞察
4. 产品策略
5. 营销策略
6. 执行计划
7. 效果评估指标
8. 风险控制措施
9. 预期效果

要求：
- 方案要具体可执行
- 包含详细的数据指标
- 考虑实际可操作性
- 语言专业且易懂
"""
        
        if requirement_data:
            prompt = f"{base_prompt}\n\n需求信息：\n"
            if requirement_data.get('company_name'):
                prompt += f"公司名称：{requirement_data['company_name']}\n"
            if requirement_data.get('industry'):
                prompt += f"行业：{requirement_data['industry']}\n"
            if requirement_data.get('target_audience'):
                prompt += f"目标受众：{requirement_data['target_audience']}\n"
            if requirement_data.get('budget_range'):
                prompt += f"预算范围：{requirement_data['budget_range']}\n"
            if requirement_data.get('campaign_goals'):
                prompt += f"营销目标：{requirement_data['campaign_goals']}\n"
            if requirement_data.get('description'):
                prompt += f"详细描述：{requirement_data['description']}\n"
        else:
            prompt = base_prompt + "\n\n请生成一份通用的营销方案模板。\n"
        
        if user_message:
            prompt += f"\n\n用户补充要求：{user_message}\n"
        
        return prompt


# 移除了复杂的连接管理和房间广播功能，符合文档要求的简单设计